import requests
import time
from openpyxl import Workbook


def query_es_data(url, headers, variable_value, wb, ws):
    # 构造带有请求头的查询ES数据的请求
    response = requests.get(url, headers=headers)
    data = response.json()

    # 获取ES返回的count值
    count = data["count"]  # 假设ES返回的数据结构中有一个名为count的字段

    # 输出查询结果到控制台
    print("Count:", count)

    # 获取当前时间并格式化
    current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())

    # 将查询时间和查询结果写入Excel表格
    row_num = ws.max_row + 1
    ws[f'A{row_num}'] = current_time
    ws[f'B{row_num}'] = count

    # 检查count是否满足条件
    if count < variable_value:
        return False, count
    else:
        return True, count


def main():
    es_url = "http://10.0.88.227:9200/api-gateway-test*/_count"  # 替换为你的Elasticsearch服务器URL
    variable_value = 10000  # 设置一个变量值，用于对比ES返回的count值

    # 设置请求头
    headers = {
        'Authorization': 'Basic ZWxhc3RpYzpjSjNAcFkyKXFHMl5tTjgk'
    }

    # 创建一个新的Excel工作簿和工作表
    wb = Workbook()
    ws = wb.active
    ws.title = "ES查询结果"
    ws['A1'] = "查询时间"
    ws['B1'] = "查询条数"

    # 记录脚本开始时间
    start_time = time.time()

    # 初始设定查询间隔为10秒
    query_interval = 20

    # 循环查询ES数据，直到count满足条件
    while True:
        success, current_count = query_es_data(es_url, headers, variable_value, wb, ws)
        if success:
            break  # 如果查询成功，则跳出循环

        time.sleep(query_interval)  # 每次根据query_interval设定的时间间隔进行查询

        # 根据差值动态调整查询间隔
        difference = variable_value - current_count
        if difference < 100:
            query_interval = 1  # 如果差值小于100，将查询间隔设定为1秒
        elif difference < 500:
            query_interval = 5  # 如果差值小于500，将查询间隔设定为5秒
    # 记录脚本结束时间
    end_time = time.time()
    duration = end_time - start_time

    # 将脚本运行时长写入Excel表格
    ws[f'A{ws.max_row + 2}'] = "入库完毕耗时"
    ws[f'B{ws.max_row}'] = f"{duration:.2f} 秒"

    # 保存Excel文件
    wb.save("../filesOutput/EsCheck.xlsx")

    print("脚本结束")


if __name__ == "__main__":
    main()
